package com.pick.base.config;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Collections;
import java.util.concurrent.TimeUnit;

/**
 * API限流配置
 * 基于Redis实现分布式限流
 *
 * @author pick
 * @since 2024-01-01
 */
@Slf4j
@Configuration
public class RateLimitConfig {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 限流Lua脚本
     * 使用滑动窗口算法实现限流
     */
    @Bean
    public RedisScript<Long> rateLimitScript() {
        String script = """
            local key = KEYS[1]
            local window = tonumber(ARGV[1])
            local limit = tonumber(ARGV[2])
            local current = tonumber(ARGV[3])
            
            -- 清除过期的记录
            redis.call('zremrangebyscore', key, 0, current - window * 1000)
            
            -- 获取当前窗口内的请求数量
            local currentCount = redis.call('zcard', key)
            
            if currentCount < limit then
                -- 添加当前请求
                redis.call('zadd', key, current, current)
                -- 设置过期时间
                redis.call('expire', key, window + 1)
                return limit - currentCount - 1
            else
                return -1
            end
            """;
        
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setScriptText(script);
        redisScript.setResultType(Long.class);
        return redisScript;
    }

    /**
     * 检查是否超过限流
     * @param key 限流键
     * @param windowSize 时间窗口大小（秒）
     * @param limit 限制次数
     * @return 剩余次数，-1表示超过限制
     */
    public long checkRateLimit(String key, int windowSize, int limit) {
        try {
            long current = System.currentTimeMillis();
            Long result = redisTemplate.execute(
                rateLimitScript(),
                Collections.singletonList(key),
                windowSize, limit, current
            );
            return result != null ? result : -1;
        } catch (Exception e) {
            log.error("限流检查异常: {}", e.getMessage(), e);
            // 异常时允许通过，避免影响正常业务
            return limit - 1;
        }
    }

    /**
     * 获取用户限流键
     * @param userId 用户ID
     * @param action 操作类型
     * @return 限流键
     */
    public String getUserRateLimitKey(String userId, String action) {
        return String.format("rate_limit:user:%s:%s", userId, action);
    }

    /**
     * 获取IP限流键
     * @param ip IP地址
     * @param action 操作类型
     * @return 限流键
     */
    public String getIpRateLimitKey(String ip, String action) {
        return String.format("rate_limit:ip:%s:%s", ip, action);
    }

    /**
     * 获取全局限流键
     * @param action 操作类型
     * @return 限流键
     */
    public String getGlobalRateLimitKey(String action) {
        return String.format("rate_limit:global:%s", action);
    }

    /**
     * 清除限流记录
     * @param key 限流键
     */
    public void clearRateLimit(String key) {
        try {
            redisTemplate.delete(key);
        } catch (Exception e) {
            log.error("清除限流记录异常: {}", e.getMessage(), e);
        }
    }

    /**
     * 获取限流剩余时间（秒）
     * @param key 限流键
     * @return 剩余时间
     */
    public long getRateLimitTtl(String key) {
        try {
            Long ttl = redisTemplate.getExpire(key, TimeUnit.SECONDS);
            return ttl != null ? ttl : 0;
        } catch (Exception e) {
            log.error("获取限流TTL异常: {}", e.getMessage(), e);
            return 0;
        }
    }

    /**
     * 预定义的限流配置
     */
    public static class RateLimitType {
        // 登录限流：每分钟最多5次
        public static final int LOGIN_WINDOW = 60;
        public static final int LOGIN_LIMIT = 5;
        
        // 注册限流：每小时最多3次
        public static final int REGISTER_WINDOW = 3600;
        public static final int REGISTER_LIMIT = 3;
        
        // 发送验证码：每分钟最多1次
        public static final int SMS_WINDOW = 60;
        public static final int SMS_LIMIT = 1;
        
        // 密码重置：每小时最多3次
        public static final int PASSWORD_RESET_WINDOW = 3600;
        public static final int PASSWORD_RESET_LIMIT = 3;
        
        // API调用：每分钟最多100次
        public static final int API_WINDOW = 60;
        public static final int API_LIMIT = 100;
        
        // 文件上传：每分钟最多10次
        public static final int UPLOAD_WINDOW = 60;
        public static final int UPLOAD_LIMIT = 10;
    }
}